import os
import yaml
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path

# 配置绝对路径
base_dir = Path.home() / "autodl-tmp/wenet/wenet/examples/senior_talk/s0/exp/conformer"
#num_epochs = 240
num_epochs = 20
# 初始化存储列表
epochs = []
losses = []
losses_att = []
losses_ctc = []
accuracies = []
lr_values = []  # 学习率变化

# 解析所有epoch文件
for n in range(num_epochs):
    file_path = base_dir / f"epoch_{n}.yaml"
    
    if not file_path.exists():
        print(f"警告: 文件 {file_path} 不存在，跳过")
        continue
    
    try:
        with open(file_path, 'r') as f:
            data = yaml.safe_load(f)
            
        # 提取训练指标
        epochs.append(n)
        losses.append(data['loss_dict']['loss'])
        losses_att.append(data['loss_dict']['loss_att'])
        losses_ctc.append(data['loss_dict']['loss_ctc'])
        accuracies.append(data['loss_dict']['acc'])
        
        # 提取学习率
        if 'lrs' in data and data['lrs']:
            lr_values.append(data['lrs'][0])
        else:
            lr_values.append(0)  # 如果缺失学习率数据，填充0
            
    except Exception as e:
        print(f"处理文件 {file_path} 时出错: {str(e)}")

# 创建图表
plt.figure(figsize=(16, 12))

# 1. 总损失曲线（对数坐标）
plt.subplot(3, 2, 1)
plt.semilogy(epochs, losses, 'b-', linewidth=2)
plt.title('Total Loss (Log Scale)')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.grid(True, linestyle='--', alpha=0.7)

# 2. 损失组件对比
plt.subplot(3, 2, 2)
plt.plot(epochs, losses, 'b-', label='Total Loss')
plt.plot(epochs, losses_att, 'r-', label='Attention Loss')
plt.plot(epochs, losses_ctc, 'g-', label='CTC Loss')
plt.title('Loss Components')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True, linestyle='--', alpha=0.7)
plt.yscale('log')

# 3. 注意力损失曲线
plt.subplot(3, 2, 3)
plt.plot(epochs, losses_att, 'r-', linewidth=2)
plt.title('Attention Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.grid(True, linestyle='--', alpha=0.7)

# 4. CTC损失曲线
plt.subplot(3, 2, 4)
plt.plot(epochs, losses_ctc, 'g-', linewidth=2)
plt.title('CTC Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.grid(True, linestyle='--', alpha=0.7)

# 5. 准确率曲线
plt.subplot(3, 2, 5)
plt.plot(epochs, accuracies, 'm-', linewidth=2)
plt.title('Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.grid(True, linestyle='--', alpha=0.7)
plt.ylim(0, max(accuracies)*1.1)  # 自动调整Y轴范围

# 6. 学习率变化
plt.subplot(3, 2, 6)
plt.plot(epochs, lr_values, 'c-', linewidth=2)
plt.title('Learning Rate Schedule')
plt.xlabel('Epoch')
plt.ylabel('Learning Rate')
plt.grid(True, linestyle='--', alpha=0.7)
plt.yscale('log')

plt.tight_layout()
plt.savefig(str(base_dir / 'training_metrics.png'), dpi=300)
print(f"图表已保存至: {base_dir / 'training_metrics.png'}")